from fastapi import Depends
from sqlmodel import select, update, desc, delete, SQLModel, and_
import typing as T
from pymilvus import MilvusClient
from ...init.Global import Session, EmbeddingSession
from ...pojo import UserKnowledgeBase, UserKnowledgeBaseVO, UserKnowledgeBaseFile


# 用户知识库的相关操作的Mapper[对外]


class _UserKnowledgeBaseMapper:
    session: Session = None
    embeddingSession: MilvusClient = None

    def __call__(self, session: Session, embeddingSession: EmbeddingSession):
        self.session = session
        self.embeddingSession = embeddingSession
        return self

    # 创建一个知识库
    async def insertByEntity(self, entity: SQLModel):
        self.session.add(entity)

    # 根据用户id获取知识库列表的功能
    async def getKnowledgeBaseListByUserId(self, userId: int):
        reuslt = (
            (
                await self.session.execute(
                    select(UserKnowledgeBase)
                    .where(UserKnowledgeBase.userId == userId)
                    .order_by(desc(UserKnowledgeBase.createTime))
                )
            )
            .scalars()
            .all()
        )
        return [UserKnowledgeBaseVO(**item.model_dump()) for item in reuslt]

    # 根据知识库id改名的功能
    async def updateKnowledgeBaseNameByKnowledgeBaseId(
        self, knowledgeBaseId: int, name: str
    ):
        await self.session.execute(
            update(UserKnowledgeBase)
            .where(UserKnowledgeBase.id == knowledgeBaseId)
            .values(alias=name)
        )

    # 根据知识库id删除知识库的功能
    async def deleteKnowledgeBaseById(
        self, knowledgeBaseId: int, collection_name: str = "default"
    ):
        await self.session.execute(
            delete(UserKnowledgeBase).where(UserKnowledgeBase.id == knowledgeBaseId)
        )
        self.embeddingSession.delete(
            collection_name, filter=f"knowledgeBaseId == {knowledgeBaseId}"
        )

    # 将知识插入知识库
    async def insertBatchToEmbedding(
        self, data: T.List[T.Dict[str, T.Any]], collection_name: str = "default"
    ):
        self.embeddingSession.insert(collection_name, data)

    # 根据知识库id搜索知识
    async def search(
        self,
        query_embeddiing: list[float],
        knowledgeBaseId: int,
        top_n: int = 1,
        collection_name: str = "default",
        type: T.Literal["text", "image"] = "text",
    ):
        result = (
            self.embeddingSession.search(
                collection_name,
                [query_embeddiing],
                filter=(
                    f"knowledgeBaseId == {knowledgeBaseId} and metadata['type'] != 'image'"
                    if type == "text"
                    else f"knowledgeBaseId == {knowledgeBaseId} and metadata['type'] == 'image'"
                ),
                output_fields=["content", "knowledgeBaseId", "metadata", "userId"],
                limit=top_n,
                anns_field="embedding",
                search_params={
                    "params": {
                        "nprobe": 16,
                    }
                },
            )
        )[0]
        return result

    # 获取图片信息
    async def getImagesInfo(
        self, knowledgeBaseId: int, count: int = -1, collection_name: str = "default"
    ):
        result = (
            self.embeddingSession.query(
                collection_name,
                filter=f"knowledgeBaseId == {knowledgeBaseId} and metadata['type'] == 'image'",
                output_fields=["content", "metadata", "userId", "knowledgeBaseId"],
            )
            if count == -1
            else self.embeddingSession.query(
                collection_name,
                filter=f"knowledgeBaseId == {knowledgeBaseId} and metadata['type'] == 'image'",
                output_fields=["content", "metadata", "userId", "knowledgeBaseId"],
                limit=count,
            )
        )
        return result

    # 根据表达式查询知识
    async def queryByExpression(self, exp: str, collection_name: str = "default"):
        result = self.embeddingSession.query(
            collection_name,
            filter=exp,
            output_fields=["content", "metadata", "userId", "knowledgeBaseId"],
        )
        return result

    # 根据id删除知识
    async def deleteEmbeddingById(self, id: int, collection_name: str = "default"):
        self.embeddingSession.delete(collection_name, ids=[id])


UserKnowledgeBaseMapper = T.Annotated[
    _UserKnowledgeBaseMapper, Depends(_UserKnowledgeBaseMapper())
]
